#include "TaaProcess.h"
#include "Utils/LogUtil.h"
#include "Render/RenderState.h"
#include "Game/GameState.h"
#include "Core/CoreGlobals.h"
#include "Game/Scene/Entity.h"
#include "Game/Component/Model.h"
#include "Render/Buffer/VertexBuffer.h"
#include "Render/Buffer/IndexBuffer.h"
#include "Render/Buffer/CustomBuffer.h"
#include "LocalShader.h"
#include "LocalScene.h"
#include "Render/VulkanTool.h"


void TaaProcess::OnInit()
{
	mLocalScene = static_cast<LocalScene*>(mGameScene);

	mDebugIndex = 3;

	mPushConstants.deltaRes = glm::vec2(mLocalScene->mDeltaWidth, mLocalScene->mDeltaWidth);
	mPushConstants.frame = GGameState->GetFrameAmount();
	mPushConstants.debugIndex = mDebugIndex;


	ENQUEUE_RENDER_COMMAND(BuildModel)([this](VkCommandBuffer commandBuffer)
		{

			RefCountPtr<RHITexture> colorImage = GetTexture("PreRenderProcess", "colorImage");
			RefCountPtr<RHITexture> normalImage = GetTexture("PreRenderProcess", "velocityImage");
			RefCountPtr<RHITexture> depthImage = GetTexture("PreRenderProcess", "depthImage");

			RefCountPtr<RHISampler> commonSampler = RHISampler::Builder()
				.SetMagFilter(VK_FILTER_NEAREST)
				.SetMinFilter(VK_FILTER_NEAREST)
				.Build();


			ColorRenderTarget::Builder colorImageBuilder;
			colorImageBuilder
				.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
				.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
				.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
			mColorImages[0] = colorImageBuilder.Build(this, "colorImage_s");
			mColorImages[1] = colorImageBuilder.Build(this, "colorImage_d");


			TextureTransfer()
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
				.SetTexture(mColorImages[1])
				.Transfer(commandBuffer);



			mDescriptor = RHIDescriptor::Builder()
				.PushShader<TaaVertexShader>()
				.PushShader<TaaFragmentShader>()
				.SetDescriptorCount(0, 2)
				.Build();

			mDescriptor->SetValue("colorImage", colorImage, commonSampler);
			mDescriptor->SetValue("velocityImage", normalImage, commonSampler);
			mDescriptor->SetValue("depthImage", depthImage, commonSampler);
			mDescriptor->SetValue("prevColorImage", mColorImages[1], commonSampler);

			mDescriptorMasks[0] = mDescriptor->AllocateDescriptor();


			mDescriptor->SetValue("colorImage", colorImage, commonSampler);
			mDescriptor->SetValue("velocityImage", normalImage, commonSampler);
			mDescriptor->SetValue("depthImage", depthImage, commonSampler);
			mDescriptor->SetValue("prevColorImage", mColorImages[0], commonSampler);

			mDescriptorMasks[1] = mDescriptor->AllocateDescriptor();


			mPipeline = GraphicsPipeline::Builder()
				.SetDescriptor(mDescriptor)
				.AddColorRenderTarget(mColorImages[0])
				.SetDstStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetDstAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.SetSrcStageMask(VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT)
				.SetSrcAccessMask(VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT)
				.Build();

			mRenderPass = RHIRenderPass::Builder()
				.AddRenderTarget(mColorImages[0])
				.AddPipeline(mPipeline)
				.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetSrcAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
				.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
				.Build();

			mFrameBuffers[0] = RHIFrameBuffer::Builder()
				.SetRenderPass(mRenderPass)
				.Build();


			mFrameBuffers[1] = RHIFrameBuffer::Builder()
				.AddRenderTarget(mColorImages[1])
				.SetRenderPass(mRenderPass)
				.Build();
		});
}

void TaaProcess::OnTick(float deltaTime)
{
	mPushConstants.deltaRes = glm::vec2(mLocalScene->mDeltaWidth, mLocalScene->mDeltaWidth);
	mPushConstants.frame = GGameState->GetFrameAmount();
	mPushConstants.debugIndex = mDebugIndex;

	ENQUEUE_RENDER_COMMAND(RecordCommand)([this](VkCommandBuffer commandBuffer)
		{

			VkRenderPassBeginInfo renderPassBeginInfo = mFrameBuffers[GRenderState->GetFrameAmount() % 2]->GetRenderPassBeginInfo();

			vkCmdBeginRenderPass(commandBuffer, &renderPassBeginInfo, VK_SUBPASS_CONTENTS_INLINE);

			vkCmdSetViewport(commandBuffer, 0, 1, &GRenderState->GetViewport());
			vkCmdSetScissor(commandBuffer, 0, 1, &GRenderState->GetScissor());

			vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipeline());

			vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS, mPipeline->GetPipelineLayout(), 0, 1, &mDescriptorMasks[GRenderState->GetFrameAmount() % 2].GetDescriptorSet(), 0, nullptr);

			vkCmdPushConstants(commandBuffer, mPipeline->GetPipelineLayout(), VK_SHADER_STAGE_FRAGMENT_BIT, 0, sizeof(PushConstants), &mPushConstants);

			vkCmdDraw(commandBuffer, 3, 1, 0, 0);

			vkCmdEndRenderPass(commandBuffer);
		});
}

void TaaProcess::OnResize()
{
	RefCountPtr<RHITexture> colorImage = GetTexture("PreRenderProcess", "colorImage");
	RefCountPtr<RHITexture> normalImage = GetTexture("PreRenderProcess", "velocityImage");
	RefCountPtr<RHITexture> depthImage = GetTexture("PreRenderProcess", "depthImage");

	RefCountPtr<RHISampler> commonSampler = RHISampler::Builder()
		.SetMagFilter(VK_FILTER_NEAREST)
		.SetMinFilter(VK_FILTER_NEAREST)
		.Build();


	ColorRenderTarget::Builder colorImageBuilder;
	colorImageBuilder
		.SetUsage(VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT)
		.SetStoreOp(VK_ATTACHMENT_STORE_OP_STORE)
		.SetFinalLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL);
	mColorImages[0] = colorImageBuilder.Build(this, "colorImage_s");
	mColorImages[1] = colorImageBuilder.Build(this, "colorImage_d");


	VkCommandBuffer onceCommandBuffer = GRenderState->BeginSingleCommand();

	TextureTransfer()
		.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
		.SetSrcAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
		.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
		.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
		.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetTexture(mColorImages[0])
		.Transfer(onceCommandBuffer);

	TextureTransfer()
		.SetSrcStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
		.SetSrcAccessMask(VK_ACCESS_SHADER_WRITE_BIT)
		.SetDstStageMask(VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT)
		.SetDstAccessMask(VK_ACCESS_SHADER_READ_BIT)
		.SetNewLayout(VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL)
		.SetTexture(mColorImages[1])
		.Transfer(onceCommandBuffer);

	GRenderState->FlushSingleCommand(onceCommandBuffer);

	mDescriptor->SetValue("colorImage", colorImage, commonSampler);
	mDescriptor->SetValue("velocityImage", normalImage, commonSampler);
	mDescriptor->SetValue("depthImage", depthImage, commonSampler);
	mDescriptor->SetValue("prevColorImage", mColorImages[1], commonSampler);

	mDescriptor->UpdateDescriptor(mDescriptorMasks[0]);


	mDescriptor->SetValue("colorImage", colorImage, commonSampler);
	mDescriptor->SetValue("velocityImage", normalImage, commonSampler);
	mDescriptor->SetValue("depthImage", depthImage, commonSampler);
	mDescriptor->SetValue("prevColorImage", mColorImages[0], commonSampler);

	mDescriptor->UpdateDescriptor(mDescriptorMasks[1]);



	mFrameBuffers[0] = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorImages[0])
		.SetRenderPass(mRenderPass)
		.Build();

	mFrameBuffers[1] = RHIFrameBuffer::Builder()
		.AddRenderTarget(mColorImages[1])
		.SetRenderPass(mRenderPass)
		.Build();
}